import java.util.ArrayList;
import java.util.HashMap;
import java.util.TreeMap;

/**
 * 
 * @author brynjar.reynisson
 */
public class StateCostMap extends TreeMap<Integer, ArrayList<StateCostWrapper>>{

	//maps in which cost list an open state is
	protected HashMap<Long, Integer> openMap = new HashMap<Long, Integer>();
	
	public void clear() {
		super.clear();
		openMap.clear();
	}
	
	public void addState(long stateId, long parentStateId, int stateCost, int heuristic) {
		addState(new StateCostWrapper(stateId, parentStateId, stateCost, heuristic));
	}
	
	public void addState(StateCostWrapper state) {
		int totalCost = state.cost+state.heuristic;
		ArrayList<StateCostWrapper> list = get(totalCost);
		if(list==null) {
			list = new ArrayList<StateCostWrapper>();
			put(totalCost, list);
		}
		list.add(state);
		openMap.put(state.stateId, totalCost);
	}

	public StateCostWrapper getCheapestState() {
		return getCheapestState(false);
	}

	public StateCostWrapper getCheapestState(boolean remove) {
		StateCostWrapper state = null;
		if(size()>0) {
			int cost = keySet().toArray(new Integer[0])[0];			
			ArrayList<StateCostWrapper> list = get(cost);
			StateCostWrapper maxState = null;
			int maxCost = 0;
			int maxIdx = -1;
			for(int i=0; i<list.size(); i++) {
				StateCostWrapper curState = list.get(i);
				if(curState.cost>=maxCost) {
					maxState = curState;
					maxCost = curState.cost;
					maxIdx = i;
				}
			}		
			state = maxState;
			if(remove) {
				openMap.remove(state.stateId);
				list.remove(maxIdx);
				if(list.size()==0) {
					remove(cost);
				}
			}
		}		
		return state;
	}
	
	public StateCostWrapper findState(long stateId) {
		StateCostWrapper state = null;
		if(openMap.containsKey(stateId)) {
			int mappedCost = openMap.get(stateId);			
			ArrayList<StateCostWrapper> l = get(mappedCost);
			for(StateCostWrapper curState: l) {
				if(curState.stateId==stateId) {
					state = curState;
					break;
				}
			}
		}		
		return state;
	}
	
	public void removeState(long stateId) {
		if(openMap.containsKey(stateId)) {
			int mappedCost = openMap.get(stateId);			
			ArrayList<StateCostWrapper> l = get(mappedCost);
			for(StateCostWrapper curState: l) {
				if(curState.stateId==stateId) {
					l.remove(curState);
					if(l.size()==0) {
						remove(mappedCost);
					}
					openMap.remove(stateId);
					break;
				}
			}
		}				
	}
	
}
